import torch
from data.dataLoader import *
from model.RF_STED import *
from model.RF_STED_No_S import *
from model.RF_STED_No_T import *
from model.RF_STED_No_RF import *
from train_model.train import *

def main():
    #获取数据
    traindata,testdata = get_od_data()
    adj1 = get_adj_matrix(ADJ_DATA_PATH +  "/adj_1.npy")
    adj2 = get_adj_matrix(ADJ_DATA_PATH +  "/adj_2.npy")
    adj3 = get_adj_matrix(ADJ_DATA_PATH +  "/adj_3.npy")
    adj4 = get_adj_matrix(ADJ_DATA_PATH +  "/adj_4.npy")
    #创建模型
    if MODEL_NAME == "RF_STED":
        model = RF_STED()
    elif MODEL_NAME == "RF_STED_No_S":
        model = RF_STED_No_S() 
    elif MODEL_NAME == "RF_STED_No_T":
        model = RF_STED_No_T() 
    elif MODEL_NAME == "RF_STED_No_RF":
        model = RF_STED_No_RF() 
    #model = ST_EN_ResCNN()    #(0.2886)
    # model = T_EN_ResCNN()     #(0.2899)
    #model3 = S_EN_ResCNN()     #(0.3023)
    #model4 = ST()              #(0.3119)
    # 加载模型
    model.load_state_dict(torch.load(PRE_MODEL_PATH))
    #训练模型
    train(model,traindata,testdata,adj1,adj2,adj3,adj4,learn_rate = 0.0001,epochs= 50)
    #保存模型
    torch.save(model.state_dict(), PRE_MODEL_PATH)
main()
